import mindspore as ms
from mindspore import nn

def cross_entropy_neg_loss(pos_score, neg_score):
    neg_score = neg_score.view(pos_score.shape[0], -1)
    score = ms.ops.concat([pos_score, neg_score], -1)
    label = ms.ops.concat([ms.ops.ones_like(pos_score), ms.ops.zeros_like(neg_score)], -1)
    return nn.BCEWithLogitsLoss()(score, label)

def hinge_loss(pos_score, neg_score):
    # an example hinge loss
    n = pos_score.shape[0]
    return (neg_score.view(n, -1) - pos_score.view(n, -1) + 1).clip_by_value(min=0).mean()


def bpr_loss_pairwise(pos_score, neg_score, epsilon=1e-9):
    bs = pos_score.shape[0]
    losses = - ms.ops.log(ms.ops.Sigmoid()(pos_score.view(bs, 1) - neg_score.view(bs, -1)) + epsilon)
    return losses.mean()


def bce_with_logits(pos_score, neg_score):
    k = neg_score.shape[0] / pos_score.shape[0]
    loss = nn.BCEWithLogitsLoss()(pos_score, ms.ops.ones_like(pos_score)) + \
           nn.BCEWithLogitsLoss()(neg_score, ms.ops.zeros_like(neg_score))
    return loss / (k + 1)


class BPRLoss(nn.LossBase):
    def construct(self, pos_score, neg_score):
        cost = []
        for pos_s, neg_s in zip(pos_score, neg_score):
            cost.append(bpr_loss_pairwise(pos_s, neg_s))
        return sum(cost)